import math
import matplotlib.pyplot as plt
import numpy as np
from scipy.spatial.transform import Rotation as Rot
from scipy.integrate import cumulative_trapezoid
from scipy.optimize import minimize_scalar

from dubins import dubins_path_planning


# 以下是新增的回旋线相关函数
def clothoid_curve(k0, k1, length, step_size):
    """
    生成回旋线（Clothoid/Euler Spiral）曲线
    参数:
    k0: 起始曲率
    k1: 结束曲率
    length: 回旋线长度
    num_points: 生成的点数
    返回:
    x, y, yaw: 回旋线上的点坐标和朝向
    """
    num_points = int(np.ceil(length/step_size))
    # 曲率变化率
    dk = (k1 - k0) / length if length > 0 else 0

    # 参数s的范围
    s = np.linspace(0, length, num_points)

    # 计算曲率随弧长的变化
    k = k0 + dk * s

    # 计算朝向角（通过积分曲率得到）
    yaw = k0 * s + 0.5 * dk * s ** 2

    # 计算坐标（通过积分cos(yaw)和sin(yaw)得到）
    dx = np.cos(yaw)
    dy = np.sin(yaw)

    # 使用 cumulative_trapezoid 替代 cumtrapz
    x = cumulative_trapezoid(dx, s, initial=0)
    y = cumulative_trapezoid(dy, s, initial=0)

    return x, y, yaw


def smooth_dubins_path(x_list, y_list, yaw_list, modes, curvature, clothoid_length_ratio, step_size):
    """
    对Dubins路径进行平滑处理，加入回旋线段

    参数:
    x_list, y_list, yaw_list: 原始路径点
    modes: 路径模式
    curvature: 曲率
    clothoid_length_ratio: 回旋线长度占圆弧长度的比例

    返回:
    smooth_x, smooth_y, smooth_yaw: 平滑后的路径
    """
    if len(x_list) < 3 or len(modes) < 2:
        return x_list, y_list, yaw_list

    # 计算路径段长度
    segment_lengths = []
    for i in range(len(modes)):
        if modes[i] == 'S':
            # 直线段长度
            segment_lengths.append(np.sqrt((x_list[-1] - x_list[0]) ** 2 +
                                           (y_list[-1] - y_list[0]) ** 2))
        else:
            # 圆弧段长度
            segment_lengths.append(abs((yaw_list[-1] - yaw_list[0]) / curvature))

    # 确定需要插入回旋线的位置
    clothoid_segments = []

    for i in range(len(modes) - 1):
        current_mode = modes[i]
        next_mode = modes[i + 1]

        # 确定当前段和下一段的曲率
        if current_mode == 'S':
            k0 = 0  # 直线曲率为0
        elif current_mode == 'L':
            k0 = curvature  # 左转圆弧曲率为正
        else:  # 'R'
            k0 = -curvature  # 右转圆弧曲率为负

        if next_mode == 'S':
            k1 = 0  # 直线曲率为0
        elif next_mode == 'L':
            k1 = curvature  # 左转圆弧曲率为正
        else:  # 'R'
            k1 = -curvature  # 右转圆弧曲率为负

        # 如果曲率有变化，需要插入回旋线
        if abs(k0 - k1) > 1e-6:
            # 计算回旋线长度
            clothoid_length = min(segment_lengths[i], segment_lengths[i + 1]) * clothoid_length_ratio

            # 生成回旋线
            clothoid_x, clothoid_y, clothoid_yaw = clothoid_curve(k0, k1, clothoid_length, step_size )

            # 确定插入位置（在路径的中间位置插入）
            insert_idx = len(x_list) // (len(modes)) * (i + 1)

            clothoid_segments.append((insert_idx, clothoid_x, clothoid_y, clothoid_yaw, k0))

    # 如果没有需要插入的回旋线，直接返回原路径
    if not clothoid_segments:
        return x_list, y_list, yaw_list

    # 构建平滑后的路径
    smooth_x, smooth_y, smooth_yaw = list(x_list), list(y_list), list(yaw_list)

    # 按插入位置从后往前处理，避免索引变化
    clothoid_segments.sort(key=lambda x: x[0], reverse=True)

    for insert_idx, c_x, c_y, c_yaw, k0 in clothoid_segments:
        # 调整回旋线的位置和方向
        offset_x = smooth_x[insert_idx] - c_x[0]
        offset_y = smooth_y[insert_idx] - c_y[0]
        offset_yaw = smooth_yaw[insert_idx] - c_yaw[0]

        adjusted_c_x = c_x + offset_x
        adjusted_c_y = c_y + offset_y
        adjusted_c_yaw = c_yaw + offset_yaw

        # 插入回旋线
        smooth_x[insert_idx:insert_idx] = adjusted_c_x.tolist()
        smooth_y[insert_idx:insert_idx] = adjusted_c_y.tolist()
        smooth_yaw[insert_idx:insert_idx] = adjusted_c_yaw.tolist()

    return smooth_x, smooth_y, smooth_yaw


def calculate_curvature(x, y):
    """
    计算路径的曲率
    """
    if len(x) < 3:
        return np.zeros(len(x))

    # 转换为numpy数组
    x = np.array(x)
    y = np.array(y)

    # 计算一阶导数
    dx = np.gradient(x)
    dy = np.gradient(y)

    # 计算二阶导数
    ddx = np.gradient(dx)
    ddy = np.gradient(dy)

    # 计算曲率
    curvature = (dx * ddy - dy * ddx) / (dx ** 2 + dy ** 2) ** 1.5

    return np.nan_to_num(curvature)


def plot_arrow(x, y, yaw, length=1.0, width=0.5, fc="r", ec="k"):
    if not isinstance(x, float):
        for (i_x, i_y, i_yaw) in zip(x, y, yaw):
            plot_arrow(i_x, i_y, i_yaw)
    else:
        plt.arrow(x, y, length * math.cos(yaw), length * math.sin(yaw), fc=fc,
                  ec=ec, head_width=width, head_length=width)
        plt.plot(x, y)


def main():
    print("Dubins path planner sample start!!")

    start_x = 0.0  # [m]
    start_y = 0.0  # [m]
    start_yaw = np.deg2rad(90.0)  # [rad]

    end_x = 5.0  # [m]
    end_y = 0.0  # [m]
    end_yaw = np.deg2rad(-90.0)  # [rad]

    radius = 8
    curvature = 1.0 / radius

    # 生成原始Dubins路径
    path_x, path_y, path_yaw, mode, lengths = dubins_path_planning(start_x,
                                                                   start_y,
                                                                   start_yaw,
                                                                   end_x,
                                                                   end_y,
                                                                   end_yaw,
                                                                   curvature)

    # 生成平滑后的Dubins路径
    smooth_x, smooth_y, smooth_yaw = smooth_dubins_path(path_x, path_y, path_yaw,
                                                        mode, curvature,
                                                        clothoid_length_ratio=0.3, step_size=0.05)

    # 计算曲率
    orig_curvature = calculate_curvature(path_x, path_y)
    smooth_curvature = calculate_curvature(smooth_x, smooth_y)

    # 绘制路径
    plt.figure(figsize=(12, 10))
    #plt.subplot(2, 1, 1)
    plt.plot(path_x, path_y, 'r-', label="original path " + "".join(mode), linewidth=2)
    plot_arrow(start_x, start_y, start_yaw)
    plot_arrow(end_x, end_y, end_yaw)
    plt.legend()
    plt.grid(True)
    plt.axis("equal")
    plt.title("Dubins Path with Clothoid Transitions")
    plt.xlabel("X [m]")
    plt.ylabel("Y [m]")
    plt.show()


    plt.figure(figsize=(12, 10))
    plt.plot(smooth_x, smooth_y, 'b-', label="smoothed path", linewidth=2)
    plot_arrow(start_x, start_y, start_yaw)
    plot_arrow(end_x, end_y, end_yaw)
    plt.legend()
    plt.grid(True)
    plt.axis("equal")
    plt.title("Dubins Path with Clothoid Transitions")
    plt.xlabel("X [m]")
    plt.ylabel("Y [m]")
    plt.show()



    # 绘制曲率
    #plt.subplot(2, 1, 2)
    plt.figure(figsize=(12, 10))
    plt.plot(orig_curvature, 'r-', label="original curvature", linewidth=2)
    #plt.plot(smooth_curvature, 'b-', label="smoothed curvature", linewidth=2)
    plt.legend()
    plt.grid(True)
    plt.title("Curvature Comparison")
    plt.xlabel("Path index")
    plt.ylabel("Curvature [1/m]")

    #plt.tight_layout()
    plt.show()


    # 绘制曲率
    #plt.subplot(2, 1, 2)
    plt.figure(figsize=(12, 10))
    #plt.plot(orig_curvature, 'r-', label="original curvature", linewidth=2)
    plt.plot(smooth_curvature, 'b-', label="smoothed curvature", linewidth=2)
    plt.legend()
    plt.grid(True)
    plt.title("Curvature Comparison")
    plt.xlabel("Path index")
    plt.ylabel("Curvature [1/m]")

    #plt.tight_layout()
    plt.show()


if __name__ == '__main__':
    main()